package com.gateway.security.endpoint;

import com.gateway.common.dto.ApiDefinition;
import com.gateway.common.enums.AuthenticationTypeEnum;
import com.gateway.common.utils.Singleton;
import com.gateway.plugin.api.util.SpringBeanUtils;
import com.gateway.security.common.authentication.client.ClientAuthenticationToken;
import com.gateway.security.common.exception.TokenAuthenticationException;
import com.gateway.security.config.SecurityConfig;
import com.gateway.security.endpoint.claim.ClaimDataCache;
import com.github.benmanes.caffeine.cache.Caffeine;
import com.github.benmanes.caffeine.cache.LoadingCache;
import com.google.common.collect.Maps;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.publisher.Mono;

import java.time.Duration;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;

/**
 * @program: api-gateway
 * @description:
 * @author: YuKai Fan
 * @create: 2025/2/10 23:51
 **/
public abstract class CacheTokenEndpointService implements TokenEndpointService {
    private static final Logger log = LoggerFactory.getLogger(CacheTokenEndpointService.class);

    @Override
    public Mono<Map<String, Object>> getClaims(ClientAuthenticationToken authenticationToken) {
        SecurityConfig securityConfig = Singleton.INST.get(SecurityConfig.class);
        if (Objects.isNull(securityConfig)) {
            return Mono.empty();
        }
        ClaimDataCache claimDataCache = ClaimDataCache.getInstance();
        Map<String, SecurityConfig.Config> configMap = securityConfig.getConfigMap();
        ApiDefinition apiDefinition = authenticationToken.getApiDefinition();
        AuthenticationTypeEnum authenticationType = apiDefinition.getAuthenticationType();
        SecurityConfig.Config config = configMap.get(authenticationType.getType());
        LoadingCache<String, Mono<Map<String, Object>>> cache = claimDataCache.obtain(authenticationType.getType());
        Duration duration = Duration.ofMillis(config.getTimeout());
        String cacheKey = authenticationToken.getToken().toString();
        return Objects.requireNonNull(cache.get(cacheKey, key -> getClaimsByToken(cacheKey, config)))
                .timeout(duration, Mono.error(new TimeoutException("Authentication took longer than timeout: " + duration)))
                .doOnError(throwable -> {
                    log.error(throwable.getMessage(), throwable);
                    cache.invalidate(cacheKey);
                    throw new TokenAuthenticationException(throwable.getMessage(), throwable);
                });
    }

    @Override
    public void refreshCache(SecurityConfig securityConfig) {
        Map<String, SecurityConfig.Config> configMap = securityConfig.getConfigMap();
        ClaimDataCache claimDataCache = ClaimDataCache.getInstance();
        Map<String, LoadingCache<String, Mono<Map<String, Object>>>> claimsCache = claimDataCache.getClaimsCache();
        if (claimsCache == null || claimsCache.isEmpty()) {
            for (Map.Entry<String, SecurityConfig.Config> entry : configMap.entrySet()) {
                SecurityConfig.Config config = entry.getValue();
                if (!config.isCached()) {
                    continue;
                }
                claimDataCache.cache(entry.getKey(), buildLoadingCache(config));
            }
        } else {
            for (String key : claimsCache.keySet()) {
                if (configMap.containsKey(key)) {
                    SecurityConfig.Config config = configMap.get(key);
                    if (!config.isCached()) {
                        claimDataCache.remove(key);
                        continue;
                    }
                    claimDataCache.cache(key, buildLoadingCache(config));
                } else {
                    claimDataCache.remove(key);
                }
            }
        }
    }

    private LoadingCache<String, Mono<Map<String, Object>>> buildLoadingCache(final SecurityConfig.Config config) {
        Caffeine<Object, Object> caffeine = Caffeine.newBuilder()
                .maximumSize(config.getMaximumSize())
                .expireAfterWrite(config.getExpireAfter(), TimeUnit.valueOf(config.getUnit()));
        if (config.isAutoRefresh()) {
            caffeine.refreshAfterWrite(
                    config.getAutoRefreshAfter(), TimeUnit.valueOf(config.getAutoRefreshUnit())
            );
        }
        return caffeine.build((key -> getClaimsByToken(key, config)));
    }

    protected abstract Mono<Map<String, Object>> getClaimsByToken(String token, SecurityConfig.Config config);
}
